import pymysql
from pymysql.cursors import DictCursor
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 数据库连接配置
DB_CONFIG = {
    'host': '150.109.6.243',
    'user': 'music_ocr_db',
    'password': 'wZktkPy4X2TpZcCY',
    'db': 'music_ocr_db',
    'charset': 'utf8mb4',
    'cursorclass': DictCursor,
}

def add_missing_columns():
    """添加缺失的数据库列"""
    try:
        # 连接数据库
        connection = pymysql.connect(**DB_CONFIG)
        cursor = connection.cursor()
        
        # 检查score_uploads表的列
        cursor.execute("SHOW COLUMNS FROM score_uploads")
        existing_columns = {column['Field']: column for column in cursor.fetchall()}
        
        # 需要添加的列
        score_columns_to_add = []
        if 'original_image_url' not in existing_columns:
            score_columns_to_add.append("ADD COLUMN original_image_url VARCHAR(255)")
        
        if 'visualization_url' not in existing_columns:
            score_columns_to_add.append("ADD COLUMN visualization_url VARCHAR(255)")
            
        if 'json_url' not in existing_columns:
            score_columns_to_add.append("ADD COLUMN json_url VARCHAR(255)")
            
        if 'cloud_storage' not in existing_columns:
            score_columns_to_add.append("ADD COLUMN cloud_storage BOOLEAN DEFAULT TRUE")
        
        # 执行score_uploads表的列添加
        if score_columns_to_add:
            alter_sql = f"ALTER TABLE score_uploads {', '.join(score_columns_to_add)}"
            logger.info(f"添加score_uploads表的列: {alter_sql}")
            cursor.execute(alter_sql)
            logger.info("成功添加score_uploads表的列")
        else:
            logger.info("score_uploads表结构已经是最新的")
        
        # 检查audio_files表的列
        cursor.execute("SHOW COLUMNS FROM audio_files")
        existing_columns = {column['Field']: column for column in cursor.fetchall()}
        
        # 需要添加的列
        audio_columns_to_add = []
        if 'midi_url' not in existing_columns:
            audio_columns_to_add.append("ADD COLUMN midi_url VARCHAR(255)")
        
        if 'wav_url' not in existing_columns:
            audio_columns_to_add.append("ADD COLUMN wav_url VARCHAR(255)")
            
        if 'cloud_storage' not in existing_columns:
            audio_columns_to_add.append("ADD COLUMN cloud_storage BOOLEAN DEFAULT TRUE")
        
        # 执行audio_files表的列添加
        if audio_columns_to_add:
            alter_sql = f"ALTER TABLE audio_files {', '.join(audio_columns_to_add)}"
            logger.info(f"添加audio_files表的列: {alter_sql}")
            cursor.execute(alter_sql)
            logger.info("成功添加audio_files表的列")
        else:
            logger.info("audio_files表结构已经是最新的")
        
        # 提交事务并关闭连接
        connection.commit()
        cursor.close()
        connection.close()
        
        logger.info("数据库更新完成")
        return True
    
    except Exception as e:
        logger.error(f"更新数据库结构时出错: {str(e)}")
        return False

if __name__ == "__main__":
    add_missing_columns() 